package com.baitiaojun.xss.core.filter;

import com.baitiaojun.xss.core.utils.JsoupUtils;
import org.apache.commons.lang3.StringUtils;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;

/**
 * servletrequest装饰器, 增加防止xss攻击能力
 */
public class ServletRequestXssWrapper extends HttpServletRequestWrapper {

    public ServletRequestXssWrapper(HttpServletRequest request) {
        super(request);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();
        Set<Map.Entry<String, String[]>> entries = parameterMap.entrySet();
        Map<String, String[]> newParameterMap = new LinkedHashMap<>();
        for (Map.Entry<String, String[]> entry : entries) {
            String[] htmlList = entry.getValue();
            for (int i = 0; i < htmlList.length; i++) {
                String html = htmlList[i];
                htmlList[i] = JsoupUtils.clean(html); //核心方法
            }
            newParameterMap.put(entry.getKey(), htmlList);
        }
        return newParameterMap;
    }

    @Override
    public String getParameter(String name) {
        String parameter = super.getParameter(name);
        if (StringUtils.isEmpty(parameter)) {
            return parameter;
        }
        return JsoupUtils.clean(parameter);
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] parameterValues = super.getParameterValues(name);
        if (parameterValues == null || parameterValues.length == 0) {
            return parameterValues;
        }
        String[] newParameterValues = new String[parameterValues.length];
        for (int i = 0; i < parameterValues.length; i++) {
            newParameterValues[i] = JsoupUtils.clean(parameterValues[i]);
        }
        return newParameterValues;
    }
}
